#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
__version__     = "0.0.3"
__copyright__   = "Copyright (C) 2023 z"
__license__     = "MIT"
__author__      = "Z"
"""

import json
import requests

from wingman import logger
from wingman.error import ErrorCode
from config.system import OPENROUTER_APP_KEY, OPENROUTER_MODEL, OPENROUTER_DEFAULT_PROMPT

class CustomAI:
    """调用openrouterAPI接口"""

    def llm(self, user_prompt):
        self.url = "https://openrouter.ai/api/v1/chat/completions"
        self.model = f"{OPENROUTER_MODEL}"
        self.default_prompt = f"{OPENROUTER_DEFAULT_PROMPT}"
        self.request_headers = {
            "Authorization": f"Bearer {OPENROUTER_APP_KEY}",
            "HTTP-Referer": "http://106.14.204.82:7000/dingtalkbot",
            "X-Title": "test"
        }
        messages = [
            {"role": "system", "content": self.default_prompt},
            {"role": "user", "content": user_prompt},
        ]
        request_json = {
            "model": self.model,
            "messages": messages,
            "max_tokens": 2048
        }
        respose = requests.request( url=self.url, method="POST", json=request_json, headers=self.request_headers )
        # 打印响应状态码和原始响应内容以进行调试
        # print("响应状态码:", respose.status_code)  # 调试输出

        # 修正:将响应内容解析为 JSON
        response_json = json.loads(respose.content.decode('utf-8'))
        # print("response_json内容:", response_json)  # 调试输出
        
        # 检查是否有错误信息
        if 'error' in response_json:
            error_msg = response_json['error']['message']
            if "Insufficient credits" in error_msg:
                raise Exception("API 余额不足，请前往 https://openrouter.ai/credits 充值")
            else:
                raise Exception(f"API 错误: {error_msg}")
        result = response_json['choices'][0]['message']['content']
        # print("调试输出:", result_msg)
        logger.info(ErrorCode.SUCCESS.value, result, __name__)
        return result

if __name__ == '__main__':
    llm = CustomAI().llm
    print( llm("你好") )
